
import json
import torch
import os
import sys
import random

from .text_diff import diff_match_patch

def get_sequence_gold_pred_tag(toks, golds, preds, step):
    
    assert len(toks) == len(golds) == len(preds), f"step:{step} {len(toks)} {len(golds)} {len(preds)}"
    
    for i, x in enumerate(golds):
        if len(golds[i]) == 3:
            golds[i] = golds[i][2]
    
    for i, x in enumerate(preds):
        if len(preds[i]) == 3:
            preds[i] = preds[i][2]
    
    token_gold_pred = [('', 'O', 'O')] + [i for i in zip(toks, golds, preds)] + [('', 'O', 'O')]
    size = len(token_gold_pred)
    raw_str = ''
    
    for k in range(1, size):
        cur = token_gold_pred[k]
        prev = token_gold_pred[k - 1]
        if prev[1:3] == cur[1:3]:
            raw_str += cur[0]
        elif prev[1:3] != cur[1:3]:
            if prev[1:3] == ('O', 'O'):
                raw_str += '' + cur[0]      #' ' + cur[0]
            elif prev[1:3] != ('O', 'O'):
                    raw_str += '/' + prev[1] + '/' + prev[2] + '' + cur[0]  #'/' + prev[1] + '/' + prev[2] + ' ' + cur[0]
    
    return raw_str.strip()


class DataProcessor(object):
    """Base class for data converters for sequence classification data sets."""
    tdiff = diff_match_patch()
    mask_tok = '[MASK]'
    rep_tok1 = '[unused97]'  #<REP>
    rep_tok2 = '[unused98]'  #</REP>
    del_tok  = '[unused99]'
    def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

    def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

    def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

    
    @classmethod
    def _read_json(self, input_file, mode='train'):
        
        if not os.path.exists(input_file):
            return []

        with open(input_file, 'r', encoding='utf8') as f:
            lines =  json.load(f)
            data = []
            
            for i, line in enumerate(lines):
                
                ori = line['original_text']
                cor = line['correct_text']
                diffs = [(0, "C") ] + self.tdiff.diff_main(ori, cor) + [(0, "S")] + [(0, "")]
                t = proc_diff(diffs, self.del_tok, self.mask_tok, self.rep_tok1, self.rep_tok2)
                
                data.append(t)
            
            return data


def proc_diff(diffs, tok_del, tok_mask, rep_tok1, rep_tok2, mode='train'):
    ori_chs = []
    cor_chs = []
    edit_tag = []
    pointing = []
    shift = 1
    
    j = 0
    while j < len(diffs) - 1:
        cur_tag, cur_txt = diffs[j]
        nxt_tag, nxt_txt = diffs[j + 1]
        
        _sz = len(cur_txt)
        _sz_nxt = len(nxt_txt)
        
        
        if cur_tag == 0:   #same text just copy
            edit_tag.extend(["S-INS|0"] * _sz)
            ori_chs.extend(cur_txt)
            cor_chs.extend(cur_txt)
            pointing.extend([shift + _ for _ in range(_sz)] )
            shift = shift + _sz
        elif cur_tag == -1:
            if nxt_tag == 1:                            #替换操作，长度可能不匹配
                if mode == 'train':
                    p = random.random()
                else:
                    p = 0   #测试一阶段
                if p <= 0.5:
                    #predict edit tag  or test step 1
                    edit_tag[-1] = f"S-INS|{len(nxt_txt)}"    #插入目标字符个数
                    edit_tag.extend(["S-DEL|1"] * _sz)        #当前字符预测为删除标签，目标词为tok_del
                    ori_chs.extend(cur_txt)
                    cor_chs.extend([tok_del] * _sz)         #目标词为tok_del，同时 预测edit_tag标签|INS|n，待纠模式
                    
                    pointing[-1] = pointing[-1] + _sz
                    pointing.extend([0] * _sz)
                    shift = shift + _sz
                
                else:
                    #FELIX MODE   test step 2
                    new_ori_part = [tok_mask] * _sz_nxt + [rep_tok1] + list(cur_txt) + [rep_tok2]
                    new_cor_part = list(nxt_txt)        + [rep_tok1] + list(cur_txt) + [rep_tok2]
                    
                    #edit_tag[-1] = f"S-INS|0"       #len(nxt_txt)             #插入目标字符个数
                    edit_tag.extend(["S-INS|0"] * len(new_ori_part))          #忽略
                    ori_chs.extend(new_ori_part)
                    cor_chs.extend(new_cor_part)                            #构造 [MASK] <REP> </REP>   FELIX MODE
                    
                    pointing[-1] = pointing[-1] +  _sz + _sz_nxt + 2
                    pointing.extend([0] * (_sz + _sz_nxt + 2))
                    shift = shift + _sz + _sz_nxt + 2
                
                j += 1                          
                
            elif nxt_tag == 0:      #删除操作
                edit_tag.extend(["S-DEL|1"] * _sz)
                ori_chs.extend(cur_txt)
                cor_chs.extend([tok_del] * _sz)     #被删除词，补充删除词占位标记
                
                pointing[-1] = pointing[-1] + _sz
                pointing.extend([0] * _sz)
                shift = shift + _sz
            
        elif cur_tag == 1:
            if nxt_tag == 0:  #纯插入操作                 
                
                if mode == 'train':
                    p = random.random()
                else:
                    p = 1
                
                if p <= 0.5 :
                    #CLOZE 完形填空 test step 2  MASK -> correct token
                    edit_tag.extend(["S-INS|0"] * len(cur_txt))       
                    ori_chs.extend([tok_mask] * len(cur_txt))
                    cor_chs.extend(cur_txt)
                    
                    pointing[-1] = pointing[-1] + _sz
                    pointing.extend([0] * _sz)
                    shift = shift + _sz
                    
                else:
                    #test step 1
                    #跳过插入字符,不带[MASK], edit_tag标签提示 INS|n
                    edit_tag[-1] = f"S-INS|{_sz}"
                    pass
        j = j + 1
    
    ori_chs[0]  = '[CLS]'
    ori_chs[-1] = '[SEP]'
    cor_chs[0]  = '[CLS]'
    cor_chs[-1] = '[SEP]'
    pointing[-1] = 0
    
    # print('-----------------------')
    # print(ori_chs)
    # print(cor_chs)
    # print(edit_tag)
    # print(pointing)
    # print('=======================')
    
    assert len(ori_chs) == len(edit_tag) == len(cor_chs) == len(pointing) , f'{len(ori_chs)} {len(edit_tag) } {len(cor_chs)} {len(pointing)}'
    return   ori_chs, edit_tag, cor_chs, pointing
    

    